package com.szl.group.config.aop;


import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.util.concurrent.RateLimiter;
import com.szl.group.common.anno.Limit;
import com.szl.group.common.enums.LimitType;
import com.szl.group.exception.BusinessException;
import com.szl.group.utils.AddressIpUtils;
import com.szl.group.utils.JoinPointUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.util.concurrent.TimeUnit;

import static com.szl.group.common.enums.ResultCodeEnum.FREQUENTLY;


/**
 * 基于Google的令牌桶算法
 */
//@Aspect
//@Configuration
//@Order(1)
public class LimitAspect {

    //根据IP分不同的令牌桶, 每天自动清理缓存
    private static LoadingCache<String, RateLimiter> caches = CacheBuilder.newBuilder()
            .maximumSize(1000)
            .expireAfterWrite(1, TimeUnit.DAYS)
            .build(new CacheLoader<String, RateLimiter>() {
                @Override
                public RateLimiter load(String key) {
                    // 新的IP初始化 每秒只发出5个令牌
                    return RateLimiter.create(5);
                }
            });

    //Service层切点  限流
    @Pointcut("@annotation(com.szl.group.common.anno.Limit)")
    public void ServiceAspect() {

    }

    @Around("ServiceAspect()")
    public Object around(ProceedingJoinPoint joinPoint) {

        Limit limitAnnotation = JoinPointUtils.getAnnotation(joinPoint, Limit.class);
        LimitType limitType = limitAnnotation.limitType();

        String key = limitAnnotation.key();
        Object obj;
        try {

            if (limitType.equals(LimitType.IP)) {
                HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
                key = AddressIpUtils.getIpAddress(request);
            }

            RateLimiter rateLimiter = caches.get(key);
            Boolean flag = rateLimiter.tryAcquire();

            if (flag) {
                obj = joinPoint.proceed();
            } else {
                throw new BusinessException(FREQUENTLY, FREQUENTLY.getMessage());
            }
        } catch (Throwable e) {
            throw new BusinessException(FREQUENTLY, FREQUENTLY.getMessage());
        }
        return obj;
    }
}
